package com.jinly.ryder.log.web;

import com.jcraft.jsch.Session;
import com.jinly.ryder.ssh.command.SshSessionUtil;
import com.jinly.ryder.ssh.command.TailCommand;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
public class RealTimeLogWebSocketHandler extends TextWebSocketHandler {

    private static ConcurrentHashMap<String, WebSocketSession> SESSION_POOL = new ConcurrentHashMap<>();
    private static ConcurrentHashMap<String, Thread> LOG_THREAD_POOL = new ConcurrentHashMap<>();

    /**
     * 建立成功事件
     *
     * @param session
     * @throws Exception
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        log.info("建立连接{}", session.getId());
        SESSION_POOL.put(session.getId(), session);
        String host = (String) session.getAttributes().get("host");
        String logFileName = (String) session.getAttributes().get("logFileName");
        Thread t = new Thread(() -> {
            log.info("开始监听日志[{}]{}", host, logFileName);
            tailLog(session, host, logFileName);
        });
        LOG_THREAD_POOL.put(session.getId(), t);
        t.start();//启动线程
    }

    private void tailLog(WebSocketSession session, String host, String logFile) {
        Session sshSession = null;
        try {
            sshSession = SshSessionUtil.connectSession(host, 22);
            TailCommand tailCommand = new TailCommand(sshSession);
            tailCommand.exec(100, logFile, (line) -> {
                try {
                    session.sendMessage(new TextMessage(line + "\n"));
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            });
        } catch (Exception e) {
            log.error("ssh log close:{}", e.getMessage(), e);
        } finally {
            SshSessionUtil.closeSession(sshSession);
            try {
                if (session != null) {
                    session.close();
                }
            } catch (Exception e) {
            }
        }
    }

    /**
     * 接收消息事件
     *
     * @param session
     * @param message
     * @throws Exception
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {

    }

    /**
     * 断开连接时
     *
     * @param session
     * @param status
     * @throws Exception
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        Thread t = LOG_THREAD_POOL.get(session.getId());
        LOG_THREAD_POOL.remove(session.getId());
        SESSION_POOL.remove(session.getId());
        t.interrupt();
        session.close();
        log.info("断开连接{}", session.getId());
    }
}
